#include "yowo_post_process.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <unordered_map>
#include <vector>

#include <algorithm>
#include <cmath>
#include <vector>

#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

using namespace std;
#define topK 40
#define NUM_CLASSES 80

const std::vector<std::vector<int>> fmp_sizes = {{28, 28}, {14, 14}, {7, 7}};
const std::vector<float> strides = {8, 16, 32};
const char *CLASS_NAMES[] = {"bend/bow(at the waist)",
                             "crawl",
                             "crouch/kneel",
                             "dance",
                             "fall down",
                             "get up",
                             "jump/leap",
                             "lie/sleep",
                             "martial art",
                             "run/jog",
                             "sit",
                             "stand",
                             "swim",
                             "walk",
                             "answer phone",
                             "brush teeth",
                             "carry/hold (an object)",
                             "catch (an object)",
                             "chop",
                             "climb (e.g. a mountain)",
                             "clink glass",
                             "close (e.g., a door, a box)",
                             "cook",
                             "cut",
                             "dig",
                             "dress/put on clothing",
                             "drink",
                             "drive (e.g., a car, a truck)",
                             "eat",
                             "enter",
                             "exit",
                             "extract",
                             "fishing",
                             "hit (an object)",
                             "kick (an object)",
                             "lift/pick up",
                             "listen (e.g., to music)",
                             "open (e.g., a window, a car door)",
                             "paint",
                             "play board game",
                             "play musical instrument",
                             "play with pets",
                             "point to (an object)",
                             "press",
                             "pull (an object)",
                             "push (an object)",
                             "put down",
                             "read",
                             "ride (e.g., a bike, a car, a horse)",
                             "row boat",
                             "sail boat",
                             "shoot",
                             "shovel",
                             "smoke",
                             "stir",
                             "take a photo",
                             "text on/look at a cellphone",
                             "throw",
                             "touch (an object)",
                             "turn (e.g., a screwdriver)",
                             "watch (e.g., TV)",
                             "work on a computer",
                             "write",
                             "fight/hit (a person)",
                             "give/serve (an object) to (a person)",
                             "grab (a person)",
                             "hand clap",
                             "hand shake",
                             "hand wave",
                             "hug (a person)",
                             "kick (a person)",
                             "kiss (a person)",
                             "lift (a person)",
                             "listen to (a person)",
                             "play with kids",
                             "push (another person)",
                             "sing to (e.g., self, a person, a group)",
                             "take (an object) from (a person)",
                             "talk to (e.g., self, a person, a group)",
                             "watch (a person)"};

struct BBox {
  float x1, y1, x2, y2;
  float score;  // 原始置信度分数
  int class_id; // 类别标签
};

struct LayerFP16 {
  const int16_t *data; // FP16数据指针
  int H;               // 特征图高度
  int W;               // 特征图宽度
  int C;               // 通道数
};

struct ComparePair {
  bool operator()(const pair<float, size_t> &a, const pair<float, size_t> &b) {
    return a.first > b.first;
  }
};

// Sigmoid 激活函数
float sigmoid(float x) { return 1.0f / (1.0f + exp(-x)); }

// 计算两个边界框的 IoU
float calculateIoU(const BBox &a, const BBox &b) {
  float interX1 = max(a.x1, b.x1);
  float interY1 = max(a.y1, b.y1);
  float interX2 = min(a.x2, b.x2);
  float interY2 = min(a.y2, b.y2);

  float interWidth = max(0.0f, interX2 - interX1);
  float interHeight = max(0.0f, interY2 - interY1);
  float interArea = interWidth * interHeight;

  float areaA = (a.x2 - a.x1) * (a.y2 - a.y1);
  float areaB = (b.x2 - b.x1) * (b.y2 - b.y1);

  return interArea / (areaA + areaB - interArea + 1e-10f);
}

// 非极大值抑制 (NMS)
vector<BBox> nms(vector<BBox> &bboxes, float nmsThresh) {
  vector<BBox> result;

  // 按分数降序排序
  sort(bboxes.begin(), bboxes.end(),
       [](const BBox &a, const BBox &b) { return a.score > b.score; });

  vector<bool> keep(bboxes.size(), true);

  for (size_t i = 0; i < bboxes.size(); ++i) {
    if (!keep[i])
      continue;
    result.push_back(bboxes[i]);
    for (size_t j = i + 1; j < bboxes.size(); ++j) {
      if (keep[j]) {
        float iou = calculateIoU(bboxes[i], bboxes[j]);
        if (iou > nmsThresh) {
          keep[j] = false;
        }
      }
    }
  }

  return result;
}

int sig[2] = {1, -1};
float result = 5.96046e-08;

float half2float(int16_t ib) {
  int16_t s, e, m;
  s = (ib >> 15) & 0x1;
  e = (ib >> 10) & 0x1f;
  m = ib & 0x3ff;

  // added by puyang.wang@lynxi.com
  {
    if (0 == e)
      return sig[s] * m * result;
    else {
      union {
        unsigned int u32;
        float f32;
      } ou;

      e = (0x1f == e) ? 0xff : (e - 15 + 127);
      ou.u32 = (s << 31) | (e << 23) | (m << 13);
      return ou.f32;
    }
  }
}

// 生成锚点
std::vector<std::pair<float, float>>
generate_anchors(const std::vector<int> &fmp_size, float stride) {
  std::vector<std::pair<float, float>> anchors;
  for (int y = 0; y < fmp_size[0]; ++y) {
    for (int x = 0; x < fmp_size[1]; ++x) {
      float cx = (x + 0.5f) * stride;
      float cy = (y + 0.5f) * stride;
      anchors.emplace_back(cx, cy);
    }
  }
  return anchors;
}
// 后处理函数
std::vector<BBox> post_process_multi_hot(
    const std::vector<LayerFP16> &conf_layers,
    const std::vector<LayerFP16> &cls_layers,
    const std::vector<LayerFP16> &reg_layers,
    const std::vector<std::vector<std::vector<float>>> &allAnchors,
    float confThresh, float nmsThresh) {

  std::vector<BBox> allBBoxes;
  for (size_t level = 0; level < 3; ++level) {
    const auto &conf_layer = conf_layers[level];
    const auto &cls_layer = cls_layers[level];
    const auto &reg_layer = reg_layers[level];
    int H = conf_layer.H;
    int W = conf_layer.W;
    float stride = strides[level];
    const auto &anchors = allAnchors[level];

    // 计算置信度分数
    std::vector<float> confScores;
    for (int h = 0; h < H; ++h) {
      for (int w = 0; w < W; ++w) {
        size_t idx = h * W * conf_layer.C + w * conf_layer.C;
        float conf = sigmoid(half2float(conf_layer.data[idx]));
        confScores.push_back(conf);
      }
    }

    // 获取TopK索引
    std::vector<std::pair<float, size_t>> scoreIndex;
    for (size_t i = 0; i < confScores.size(); ++i) {
      scoreIndex.emplace_back(confScores[i], i);
    }
    std::sort(scoreIndex.rbegin(), scoreIndex.rend()); // 降序排序

    int k = std::min(topK, static_cast<int>(scoreIndex.size()));
    std::vector<size_t> topkIndices;
    for (int i = 0; i < k; ++i) {
      if (scoreIndex[i].first < confThresh)
        break;
      topkIndices.push_back(scoreIndex[i].second);
    }

    // 处理每个候选框
    for (const auto &idx : topkIndices) {
      int h = idx / reg_layer.W; // 使用实际宽度
      int w = idx % reg_layer.W;
      const auto &anchor = anchors[idx];

      // 解码回归参数
      size_t reg_idx = (h * reg_layer.W + w) * reg_layer.C;
      float dx = half2float(reg_layer.data[reg_idx]);
      float dy = half2float(reg_layer.data[reg_idx + 1]);
      float dw = half2float(reg_layer.data[reg_idx + 2]);
      float dh = half2float(reg_layer.data[reg_idx + 3]);

      float ctrX = anchor[0] + dx * stride;
      float ctrY = anchor[1] + dy * stride;
      float width = std::exp(dw) * stride;
      float height = std::exp(dh) * stride;

      BBox bbox;
      bbox.x1 = ctrX - width * 0.5f;
      bbox.y1 = ctrY - height * 0.5f;
      bbox.x2 = ctrX + width * 0.5f;
      bbox.y2 = ctrY + height * 0.5f;

      // 处理类别
      size_t cls_idx = h * W * cls_layer.C + w * cls_layer.C;
      float max_score = 0;
      int cls_id = 0;
      for (int c = 0; c < NUM_CLASSES; ++c) {
        float score = sigmoid(half2float(cls_layer.data[cls_idx + c]));
        if (score > max_score) {
          max_score = score;
          cls_id = c;
        }
      }

      bbox.score = std::sqrt(confScores[idx] * max_score);
      bbox.class_id = cls_id;
      allBBoxes.push_back(bbox);
    }
  }

  // 执行NMS
  return nms(allBBoxes, nmsThresh);
}

// 主处理函数
int lynYowoPostProcess(YowoPostInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;

#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif

  int img_w = post_info->img_width;
  int img_h = post_info->img_height;

  std::vector<LayerFP16> conf_layers, cls_layers, reg_layers;
  int16_t *fp16_base = static_cast<int16_t *>(output_tensor);
  size_t current_offset = 0;

  for (int i = 0; i < post_info->model_output_len; ++i) {
    int level = i % 3; // 确定特征层级别
    int H = fmp_sizes[level][0];
    int W = fmp_sizes[level][1];
    int M = H * W;
    int C = (i < 3) ? 1 : (i < 6) ? NUM_CLASSES : 4;

    LayerFP16 layer;
    layer.data = fp16_base + current_offset;
    layer.H = H;
    layer.W = W;
    layer.C = C;

    // 分组存储
    if (i < 3)
      conf_layers.push_back(layer);
    else if (i < 6)
      cls_layers.push_back(layer);
    else
      reg_layers.push_back(layer);

    current_offset += M * C;
  }

  static std::vector<std::vector<std::vector<float>>> all_anchors;
  static bool anchors_initialized = false;
  if (!anchors_initialized) {
    // 生成锚点
    for (size_t level = 0; level < 3; ++level) {
      auto anchors = generate_anchors(
          {conf_layers[level].H, conf_layers[level].W}, strides[level]);
      std::vector<std::vector<float>> level_anchors;
      for (const auto &a : anchors) {
        level_anchors.push_back({a.first, a.second});
      }
      all_anchors.push_back(level_anchors);
    }
    anchors_initialized = true;
  }

  // 执行后处理
  auto all_bboxes = post_process_multi_hot(
      conf_layers, cls_layers, reg_layers, all_anchors,
      post_info->conf_threshold, post_info->nms_threshold);

  // 填充输出结构
  uint16_t count = 0;
  auto clamp_normalize = [post_info](float coord, uint32_t max_dim) {
    float normalized =
        coord / std::max(post_info->model_input_h, post_info->model_input_w);
    normalized =
        (normalized < 0.0f) ? 0.0f : (normalized > 1.0f) ? 1.0f : normalized;

    uint32_t pixel = static_cast<uint32_t>(std::round(normalized * max_dim));
    return pixel >= max_dim ? max_dim - 1 : pixel; // 确保不越界
  };

  for (const auto &box : all_bboxes) {
    if (box.score < post_info->vis_threshold)
      continue;
    if (count >= BOX_MAX_NUM)
      break;

    lynBox out_box{.xmin = clamp_normalize(box.x1, img_w),
                   .ymin = clamp_normalize(box.y1, img_h),
                   .xmax = clamp_normalize(box.x2, img_w),
                   .ymax = clamp_normalize(box.y2, img_h),
                   .score = box.score,
                   .id = box.class_id};

    strncpy(out_box.label, CLASS_NAMES[box.class_id], TEXT_MAX_LEN - 1);
    out_box.label[TEXT_MAX_LEN - 1] = '\0';

    boxesInfo->boxes[count++] = out_box;
  }

  boxesInfo->boxesNum = count;

  return 0;
}